from fastapi import APIRouter, UploadFile, File, Form
import torch
from torchvision import models, transforms
from PIL import Image
from core.config import settings  # 新添加这一行
from schemas import PredictResponse


router = APIRouter(prefix="/inference", tags=["inference"])

@router.post("/predict", response_model=PredictResponse)
async def predict(image: UploadFile = File(...), model_id: str = Form(...)):
    # 加载模型（示例使用ResNet）
    model = models.resnet50(pretrained=False)

    # 加载状态字典，并剥离 'module.' 前缀
    state_dict = torch.load(f"{settings.MODELS_DIR}/{model_id}", map_location="cpu")
    new_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith("module."):
            new_state_dict[k[7:]] = v  # 移除 'module.' (7 个字符)
        else:
            new_state_dict[k] = v
    model.load_state_dict(new_state_dict)  # 现在加载修改后的字典
    model.eval()

    # 图像预处理
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    img = Image.open(image.file)
    img_t = transform(img)
    batch_t = torch.unsqueeze(img_t, 0)

    # 推理
    with torch.no_grad():
        out = model(batch_t)
    probabilities = torch.nn.functional.softmax(out[0], dim=0)
    # 假设标签为wg01, wg02等；替换为实际
    predictions = [
        {"label": "wg01", "probability": probabilities[0].item()},
        {"label": "wg02", "probability": probabilities[1].item()}
    ]
    return PredictResponse(success=True, predictions=predictions)


